# -*- coding: utf-8 -*-
"""
Module for creating publication-quality visualizations of simulation results.

This module uses PyVista to automate the generation of scenes similar to
those described in the visualization_guide.md.
"""
import os
import logging
import pyvista as pv
import numpy as np

class Visualizer:
    """
    Handles the creation of visualizations from VTK files.
    """
    def __init__(self, config):
        """
        Initializes the Visualizer.

        Args:
            config: The simulation configuration object.
        """
        self.config = config
        self.output_dir = self.config.VISUALIZATION_OUTPUT_DIR
        
        # Ensure the output directory exists
        if not os.path.exists(self.output_dir):
            os.makedirs(self.output_dir)
            logging.info(f"Created visualization output directory: {self.output_dir}")
        
        # Set PyVista to run in off-screen mode for automated execution
        pv.set_plot_theme("document")
        pv.global_theme.background = 'white'
        pv.global_theme.font.color = 'black'

    def create_scene(self, vtk_filepath: str, time_step_index: int):
        """
        Generates and saves a slice-focused scene from a given VTK file.

        Optimized version that shows internal structure efficiently:
        - Single primary slice through center
        - Minimal computational overhead
        - Clear view of internal mesh structure

        Args:
            vtk_filepath (str): Path to the input VTK file.
            time_step_index (int): The current time step index, used for the output filename.
        """
        if not os.path.exists(vtk_filepath):
            logging.error(f"VTK file not found: {vtk_filepath}")
            return

        try:
            # 1. Load the dataset
            mesh = pv.read(vtk_filepath)
            logging.debug(f"Loaded mesh with {mesh.n_points} points, {mesh.n_cells} cells")

            # 2. Create a plotter object with smaller window for faster rendering
            plotter = pv.Plotter(off_screen=True, window_size=[1200, 800])
            
            # Add title
            plotter.add_text(f"Dragonfly CFD - Time Step {time_step_index}", 
                           position='upper_left', font_size=14, color='black')

            # 3. Add outer surface as very transparent boundary (minimal processing)
            surface = mesh.extract_surface()
            plotter.add_mesh(surface, color='lightgray', opacity=0.08, 
                           show_edges=False)

            # 4. Create primary slices to show internal structure
            center = mesh.center
            
            # XY slice (horizontal) for pressure visualization
            try:
                xy_slice = mesh.slice(normal='z', origin=center)
                if xy_slice.n_cells > 0:
                    if "pressure" in mesh.point_data:
                        plotter.add_mesh(xy_slice, scalars="pressure", cmap="RdBu_r", 
                                       opacity=0.9, show_edges=True, 
                                       scalar_bar_args={'title': 'Pressure (Pa)', 'position_x': 0.85})
                    else:
                        plotter.add_mesh(xy_slice, color="lightblue", opacity=0.8, 
                                       show_edges=True)
                    logging.debug("Added XY slice for pressure")
            except Exception as e:
                logging.warning(f"Could not create XY slice: {e}")

            # XZ slice (vertical) for velocity magnitude visualization
            try:
                xz_slice = mesh.slice(normal='y', origin=center)
                if xz_slice.n_cells > 0:
                    if "velocity" in mesh.point_data:
                        # Calculate velocity magnitude for the slice
                        velocity_data = mesh.point_data["velocity"]
                        if len(velocity_data.shape) > 1:
                            # Calculate velocity magnitude
                            velocity_mag = np.linalg.norm(velocity_data, axis=1)
                            
                            # Map velocity magnitude to slice points
                            # Get original point IDs for the slice
                            if "vtkOriginalPointIds" in xz_slice.point_data:
                                orig_ids = xz_slice.point_data["vtkOriginalPointIds"]
                                slice_vel_mag = velocity_mag[orig_ids]
                            else:
                                # Fallback: use slice point indices directly
                                slice_vel_mag = velocity_mag[:xz_slice.n_points]
                            
                            xz_slice["velocity_magnitude"] = slice_vel_mag
                            plotter.add_mesh(xz_slice, scalars="velocity_magnitude", cmap="viridis", 
                                           opacity=0.8, show_edges=True,
                                           scalar_bar_args={'title': 'Velocity Magnitude (m/s)', 'position_x': 0.15})
                            logging.debug("Added XZ slice for velocity magnitude")
                        else:
                            plotter.add_mesh(xz_slice, color="lightcoral", opacity=0.7, 
                                           show_edges=True)
                            logging.debug("Added XZ slice (no velocity data)")
                    else:
                        plotter.add_mesh(xz_slice, color="lightcoral", opacity=0.7, 
                                       show_edges=True)
                        logging.debug("Added XZ slice (no velocity field)")
            except Exception as e:
                logging.warning(f"Could not create XZ slice: {e}")

            # 5. Add velocity vectors (simplified - only if data exists and mesh is small enough)
            if "velocity" in mesh.point_data and mesh.n_points < 5000:
                try:
                    # Simple subsampling without complex filtering
                    n_skip = max(1, mesh.n_points // 100)  # Target ~100 glyphs
                    indices = list(range(0, mesh.n_points, n_skip))
                    sampled_points = mesh.points[indices]
                    sampled_velocity = mesh["velocity"][indices]
                    
                    # Create a polydata object with sampled points
                    sampled_mesh = pv.PolyData(sampled_points)
                    sampled_mesh["velocity"] = sampled_velocity
                    
                    # Create arrow glyphs for velocity vectors
                    glyphs = sampled_mesh.glyph(
                        orient="velocity", 
                        scale="velocity", 
                        factor=0.05,  # Smaller arrows
                        geom=pv.Arrow()
                    )
                    plotter.add_mesh(glyphs, color="red", opacity=0.7)
                    logging.debug("Added velocity vectors")
                except Exception as e:
                    logging.warning(f"Could not add velocity vectors: {e}")

            # 6. Set camera position for optimal slice viewing
            plotter.camera_position = 'xy'
            plotter.camera.azimuth = 30
            plotter.camera.elevation = 20
            plotter.camera.zoom(1.3)

            # 7. Add minimal annotation
            plotter.add_text("Internal Slice View\nHorizontal: Pressure (Red-Blue) | Vertical: Velocity Magnitude (Viridis)", 
                           position='lower_left', font_size=10, color='black')

            # 8. Save the output
            output_filename = os.path.join(self.output_dir, f"scene_{time_step_index:05d}.png")
            plotter.screenshot(output_filename)
            logging.debug(f"Saved slice visualization to {output_filename}")
            
            plotter.close()

        except Exception as e:
            logging.error(f"Failed to create visualization for {vtk_filepath}: {e}") 